Skip to content

[RFC][ROCm][AITER] Keep all AITER kernels in _aiter_ops class like _custom_ops and _ipex_ops#24490

Merged
tjtanaa merged 60 commits intovllm-project:mainfrom
EmbeddedLLM:rfc-aiter-ops
Nov 10, 2025
Merged

[RFC][ROCm][AITER] Keep all AITER kernels in _aiter_ops class like _custom_ops and _ipex_ops#24490
tjtanaa merged 60 commits intovllm-project:mainfrom
EmbeddedLLM:rfc-aiter-ops

Conversation

@vllmellm
Copy link
Copy Markdown
Contributor

@vllmellm vllmellm commented Sep 9, 2025 •
edited by github-actions bot
Loading

Purpose

This PR introduces _aiter_ops.py as proposed in the RFC here. The aiter_ops namespace provides several key benefits:

  • Centralized kernel registration: Ensures that kernels from the aiter package are properly registered

  • Environment availability checks: Encapsulates aiter support detection and environment compatibility validation

  • Reduced code duplication: Eliminates the need for duplicate helper functions, namely checking device compability and environment varible enability checks across different vLLM modules.

This implementation establishes the foundation for future refactoring efforts, where existing kernels throughout the vLLM repository will be migrated to use this unified approach for better maintainability and consistency.

This PR uses 5ee37dce commit from aiter repo.

Test Plan

Test models that are afftected by this change, using lm_eval on gsm8k dataset.

environment setting

Step 1: run vllm serve

VLLM_USE_V1=1 \ VLLM_ROCM_USE_AITER=1 \ SAFETENSORS_FAST_GPU=1 \ VLLM_DISABLE_COMPILE_CACHE=1 \ vllm serve $MODEL_NAME --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE", "cudagraph_capture_sizes": [1,2,4,8,16,24,32]}' --trust-remote-code --swap-space 16 --distributed-executor-backend mp

Step 2: run lm_eval

lm_eval --model local-completions --tasks gsm8k --model_args model=$MODEL_NAME,base_url=http://localhost:8000/v1/completions --trust_remote_code --num_fewshot 5 --batch_size 256

Test Results

deepseek-ai/DeepSeek-V3 -tp 8 --block-size 1 --max-model-len 32768 --max_seq_len_to_capture 32768

Tasks Version Filter n-shot Metric Value Stderr
gsm8k 3 flexible-extract 5 exact_match ↑ 0.9500 ± 0.006
strict-match 5 exact_match ↑ 0.9507 ± 0.006

meta-llama/Llama-4-Scout-17B-16E-Instruct -tp 8 --max-model-len 8192

Tasks Version Filter n-shot Metric Value Stderr
gsm8k 3 flexible-extract 5 exact_match ↑ 0.9174 ± 0.0076
strict-match 5 exact_match ↑ 0.8999 ± 0.0083

meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8 -tp 8 --max-model-len 8192

Tasks Version Filter n-shot Metric Value Stderr
gsm8k 3 flexible-extract 5 exact_match ↑ 0.9272 ± 0.0072
strict-match 5 exact_match ↑ 0.9287 ± 0.0071

serve Qwen/Qwen3-235B-A22B-FP8 -tp 4

Tasks Version Filter n-shot Metric Value Stderr
gsm8k 3 flexible-extract 5 exact_match ↑ 0.8704 ± 0.0093
strict-match 5 exact_match ↑ 0.8635 ± 0.0095

mistralai/Mixtral-8x7B-Instruct-v0.1 -tp 2

Tasks Version Filter n-shot Metric Value Stderr
gsm8k 3 flexible-extract 5 exact_match ↑ 0.6452 ± 0.0132
strict-match 5 exact_match ↑ 0.6422 ± 0.0132

mistralai/Mixtral-8x7B-Instruct-v0.1 -tp 2 --quantization fp8

Tasks Version Filter n-shot Metric Value Stderr
gsm8k 3 flexible-extract 5 exact_match ↑ 0.6058 ± 0.0135
strict-match 5 exact_match ↑ 0.6027 ± 0.0135

meta-llama/Meta-Llama-3.3-70B-Instruct -tp 2

Tasks Version Filter n-shot Metric Value Stderr
gsm8k 3 flexible-extract 5 exact_match ↑ 0.9386 ± 0.0066
strict-match 5 exact_match ↑ 0.9128 ± 0.0078

meta-llama/Meta-Llama-3.3-70B-Instruct -tp 2 --quantization fp8

Tasks Version Filter n-shot Metric Value Stderr
gsm8k 3 flexible-extract 5 exact_match ↑ 0.9424 ± 0.0064
strict-match 5 exact_match ↑ 0.9128 ± 0.0078

amd/Llama-3.3-70B-Instruct-FP8-KV -tp 2

Tasks Version Filter n-shot Metric Value Stderr
gsm8k 3 flexible-extract 5 exact_match ↑ 0.9378 ± 0.0067
strict-match 5 exact_match ↑ 0.9030 ± 0.0082
Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

vllmellm added 12 commits August 28, 2025 12:18
@vllmellm
move all aiter kernel ops registration to _aiter_ops.py and centraliz…
89fc090 
…e the aiter ops enability and support

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@vllmellm
Merge remote-tracking branch 'origin/main' into rfc-aiter-ops
2f07bcc
@vllmellm
clean remaining of aiter ops
b67838b 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@vllmellm
fix circular import
9c1387b 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@vllmellm
avoid circular import
414c638 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@vllmellm
Merge remote-tracking branch 'origin/main' into rfc-aiter-ops
ad0d26c
@vllmellm
fix fused_topk return
e878b9a 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@vllmellm
add missing aiter gemm_w8a8_block scale kernel in iater ops.
c9663f2 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@vllmellm
fix ops name
9cfcaf4 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@vllmellm
Merge remote-tracking branch 'origin/main' into rfc-aiter-ops
ccf5a0a
@vllmellm
remove unnecessary cache decorator
6ffe492 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@vllmellm
Merge remote-tracking branch 'origin/main' into rfc-aiter-ops
ba3fa63 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@mergify mergify bot added rocm Related to AMD ROCm v1 labels Sep 9, 2025
@vllmellm
fix precommit error
9c8a7c7 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@mergify
Copy link
Copy Markdown

mergify bot commented Sep 10, 2025

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @vllmellm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Sep 10, 2025
@vllmellm
Merge remote-tracking branch 'origin/main' into rfc-aiter-ops
bdd7b74 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@mergify mergify bot removed the needs-rebase label Sep 12, 2025
@vllmellm vllmellm marked this pull request as ready for review September 12, 2025 21:10
@mergify mergify bot added the needs-rebase label Nov 5, 2025
@vllmellm
Merge remote-tracking branch 'origin/main' into rfc-aiter-ops
c2b050d 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@mergify mergify bot removed the needs-rebase label Nov 5, 2025
@mergify
Copy link
Copy Markdown

mergify bot commented Nov 7, 2025

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @vllmellm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Nov 7, 2025
@vllmellm
Merge remote-tracking branch 'origin/main' into rfc-aiter-ops
488b5b8 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@mergify mergify bot removed the needs-rebase label Nov 7, 2025
@DarkLight1337 DarkLight1337 added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 7, 2025
@tjtanaa tjtanaa enabled auto-merge (squash) November 9, 2025 16:36
@tjtanaa tjtanaa merged commit f080a83 into vllm-project:main Nov 10, 2025
67 checks passed
@sarckk sarckk mentioned this pull request Nov 10, 2025
Merged
5 tasks
@mgoin mgoin mentioned this pull request Nov 10, 2025
Merged
5 tasks
@mgoin
Copy link
Copy Markdown
Member

mgoin commented Nov 10, 2025

@tjtanaa @vllmellm it seems this PR introduced mandatory imports of rocm.py, showing

WARNING 11-10 22:31:06 [rocm.py:39] Failed to import from amdsmi with ModuleNotFoundError("No module named 'amdsmi'")
WARNING 11-10 22:31:06 [rocm.py:50] Failed to import from vllm._rocm_C with ModuleNotFoundError("No module named 'vllm._rocm_C'")

by default on nvidia devices. Can we move these to be lazy imports? I opened #28428 for fp8_utils.py but then saw that this is applied to more files

This was referenced Nov 11, 2025
Merged
Open
@hurricane-dorian hurricane-dorian mentioned this pull request Nov 15, 2025
Open
devpatelio pushed a commit to SumanthRH/vllm that referenced this pull request Nov 29, 2025
@vllmellm @ProExpertProg @devpatelio
[RFC][ROCm][AITER] Keep all AITER kernels in _aiter_ops class like …
e4cbdb6 
…`_custom_ops` and `_ipex_ops` (vllm-project#24490)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
@tjtanaa tjtanaa mentioned this pull request Dec 18, 2025
Closed
1 task
@Klaud-Cold Klaud-Cold mentioned this pull request Feb 26, 2026
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ProExpertProg ProExpertProg ProExpertProg left review comments

@tjtanaa tjtanaa tjtanaa approved these changes

@gshtras gshtras Awaiting requested review from gshtras

@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner

@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 is a code owner

@tdoublep tdoublep Awaiting requested review from tdoublep

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat is a code owner

@njhill njhill Awaiting requested review from njhill

@ywang96 ywang96 Awaiting requested review from ywang96

@comaniac comaniac Awaiting requested review from comaniac

@alexm-redhat alexm-redhat Awaiting requested review from alexm-redhat

@mgoin mgoin Awaiting requested review from mgoin mgoin is a code owner

@LucasWilkinson LucasWilkinson Awaiting requested review from LucasWilkinson

@pavanimajety pavanimajety Awaiting requested review from pavanimajety pavanimajety is a code owner

+1 more reviewer

@SageMoore SageMoore SageMoore approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

deepseek Related to DeepSeek models documentation Improvements or additions to documentation ready ONLY add when PR is ready to merge/full CI is needed rocm Related to AMD ROCm v1

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@vllmellm @ProExpertProg @tjtanaa @mgoin @SageMoore @DarkLight1337